# Tests Multi-variate Normal and Wishart Distributions

import math
from testmain import MyTestCase, testmodule

from pyblog import *
from pyblog.misc.poly import Polynomial
from random import Random

try:
  import numpy as N
  import numpy.linalg as L
except:
  print "NumPy not installed: skipping Multivariate Normal"
  raise

def scatter(x): return x * x.transpose()

class test_mnormal_dist(MyTestCase):
  def testmeancov(self):
    l = N.matrix([[1,0,0],[2,3,4],[5,6,7]])
    sigma = l * l.transpose()
    mean = N.matrix([1, 2, 3]).transpose()
    dist = MNormal(mean, L.inv(sigma))
    rand = Random()
    n = 10000
    samplemean = sum(dist.sample(rand)/n for i in range(n))
    samplecov = sum(scatter(dist.sample(rand) - mean)/n for i in range(n))
    meanerr = float((mean - samplemean).transpose() * (mean - samplemean))
    coverr = L.det(samplecov - sigma)
    
    if dispstats:
      print "True Mean", str(mean.transpose())
      print "Sample Mean", str(samplemean.transpose())
      print "Mean Error", meanerr
      print "True Cov\n", str(sigma)
      print "Sample Cov\n", str(samplecov)
      print "Cov Err\n", coverr
      
    self.assertAbsEqual(meanerr, 0, .5, "sample mean")
    self.assertAbsEqual(coverr, 0, .001, "sample covariance")

  def testprob(self):
    dist = MNormal([2,3], [[2,1],[1,2]])
    totalprob = 0
    for i in range(-100, 100, 1):
      for j in range(-100, 100, 1):
        x = N.matrix([i * .1, j * .1]).transpose()
        totalprob += dist.prob(x) * .1 * .1
    if dispstats:
      print "Total Prob", totalprob
    self.assertAbsEqual(totalprob, 1, .1, "Total Probability")

class test_wishart_dist(MyTestCase):
  def testmean(self):
    cov = N.matrix([[1,2,5],[2, 29, 56], [5, 56, 110]])
    dist = Wishart(L.inv(cov), 4)
    rand = Random()
    n = 10000
    samplemean = sum(dist.sample(rand)/n for i in range(n))
    meanerr = L.det(4 * cov - samplemean)
    if dispstats:
      print "True Mean:\n", str(4 * cov)
      print "Sample Mean:\n", str(samplemean)
      print "Mean Error:", meanerr
    self.assertAbsEqual(meanerr, 0, .01, "Sample Mean")

  def testprob(self):
    # the probability values are not normalized for the Wishart distribution
    # hence we need to test it differently.
    cov = N.matrix([[.1, .05], [.05, .1]])
    dist = Wishart(L.inv(cov), 4)
    rand = Random()
    compmean = N.zeros((2,2))
    totalwt = 0.0
    for i in range(-10, 10, 1):
      for j in range(-10, 10, 1):
        for k in range(-10, 10, 1):
          x = N.matrix([[i*.1, j*.1], [j*.1, k*.1]])
          wt = dist.prob(x)
          compmean += x * wt
          totalwt += wt
    compmean /= totalwt
    truemean = 4 * cov
    meanerr = L.det(compmean - truemean)
    if dispstats:
      print "True Mean:\n", str(truemean)
      print "Computed Mean:\n", str(compmean)
      print "Mean Error:", meanerr
    self.assertAbsEqual(meanerr, 0, .001, "mean error")

class testConjugacy(MyTestCase):
  def testmean(self):
    x = Polynomial(N.matrix([1,2]).transpose())
    y = N.matrix([2, 3]).transpose()
    z = N.matrix([[2, 1], [1, 2]])
    s0 = MNormal.SufficientStats()
    s0.add(y)
    l0 = MNormal(x, z)(s0)

    l0 = l0 * 2

    l0 = .5 * l0

    if dispstats:
      print "Multivariate mean likelihood:", l0

    self.assertTrue((l0.get_mean_invcov()[0] == y).all(),
                    "MNormal mean likelihood mean")
    self.assertTrue((l0.get_mean_invcov()[1] == z).all(),
                    "MNormal mean likelihood invcov")

    y0 = N.matrix([1,1]).transpose()
    z0 = N.matrix([[.5, .2], [.2, .3]])
    post = l0 * MNormal(y0, z0)

    if dispstats:
      print "Posterior:", post

    self.assertTrue((post.get_mean_invcov()[0] \
                     == L.inv(z + z0) * (z*y + z0*y0)).all(),
                    "MNormal posterior mean")
    self.assertTrue((post.get_mean_invcov()[1] \
                     == (z + z0)).all(),
                    "MNormal posterior invcov")


    
  def testinvcov(self):
    x = Polynomial(N.matrix([[6, -1],[-1, 5]]).transpose())
    
    s0 = MNormal.SufficientStats()
    s0.add(N.matrix([1,2]).transpose())
    l0 = MNormal([1,3], x)(s0)
    
    l0 = l0 * 2
    
    l0 = 2 * l0

    if dispstats:
      print "Multivariate invcov likelihood:", str(l0)

    self.assertTrue((l0.get_wishart_invscale_dof()[0] \
                     == N.matrix([[0,0],[0,1]])).all(),
                    "MNormal invcov likelihood -- invscale")
    self.assertTrue(l0.get_wishart_invscale_dof()[1] == 4,
                    "MNormal invcov likelihood -- dof")
    
    post = l0 * Wishart(N.matrix([[3, -2], [-2, 4]]), 5)
    
    if dispstats:
      print "Posterior:", str(post)

    self.assertTrue((post.get_wishart_invscale_dof()[0] \
                     == N.matrix([[3,-2],[-2,5]])).all(),
                    "Wishart posterior -- invscale")
    self.assertTrue(post.get_wishart_invscale_dof()[1] == 6,
                    "Wishart posterior -- dof")

    
if __name__ == "__main__":
  testmodule(__import__("testmnormal"))
